import os,cv2
import numpy as np
from math import ceil

from dataset.data_aug import *
from dataset.dataset import collate_fn
import time
import torch
import torch.utils.data as torchdata
import torch.utils.data as data
from torch.autograd import Variable
# from torchvision import transforms

# from classify3.resnet import *
from models.MobileNet import MobileNet
from models.resnet import resnet50


class dataset_pred(data.Dataset):
    def __init__(self, imgs,label, transforms=None):
        self.imgs = imgs
        self.labels = label
        self.transforms = transforms

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, item):
        # img_path = self.paths[item]
        img =self.imgs[0]
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        if self.transforms is not None:
            img = self.transforms(img)
        label = self.labels[item]
        return torch.from_numpy(img).float(), label


def class5_init(resume):


    test_transforms= Compose([
            ExpandBorder(size=(272,272),resize=True),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    # resnet
    # model =resnet50(pretrained=True)
    # model.avgpool = torch.nn.AdaptiveAvgPool2d(output_size=1)
    # model.fc = torch.nn.Linear(model.fc.in_features,4)

    #mobilenet
    model = MobileNet(5, alpha=0.5)
    # model.fc = torch.nn.Conv2d(in_channels=512, out_channels=5, kernel_size=1, padding=0)
    model.fc = torch.nn.Linear(model.fc.in_features, 5)

    model.load_state_dict(torch.load(resume),strict=True)
    model = model.cuda()
    model.eval()
    return model,test_transforms


def class2_init(resume):

    test_transforms= Compose([
            # ExpandBorder(size=(48,26),resize=True),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    # resnet
    # model =resnet50(pretrained=True)
    # model.avgpool = torch.nn.AdaptiveAvgPool2d(output_size=1)
    # model.fc = torch.nn.Linear(model.fc.in_features, 2)

    # mobilenet
    model = MobileNet(2, alpha=0.5)
    # model.fc = torch.nn.Conv2d(in_channels=512, out_channels=5, kernel_size=1, padding=0)
    model.fc = torch.nn.Linear(model.fc.in_features, 3)

    model.load_state_dict(torch.load(resume), strict=True)
    model = model.cuda()
    model.eval()
    return model,test_transforms

def class5_forward(testimage,model,test_transforms):

    ## 4  其他
    # 3 竖直
    # 2水平倾斜
    # 1 亮
    # 0 灭
    start =time.time()

    input = test_transforms(cv2.cvtColor(testimage[0],cv2.COLOR_BGR2RGB))
    input = input[np.newaxis, :]
    input = torch.FloatTensor(input)
    input = input.to(torch.device("cuda"))
    output = model(input)

    if isinstance(output, list):
        output = (output[0] + output[1]) / 2

    props, preds = torch.max(output, 1)
    test_preds = preds.data.cpu().numpy()[0]
    return test_preds


def class2_forward(testimage,model,test_transforms):

    ## 4  其他
    # 3 竖直
    # 2水平倾斜
    # 1 亮
    # 0 灭
    start =time.time()

    input = test_transforms(cv2.cvtColor(testimage[0],cv2.COLOR_BGR2RGB))
    input = input[np.newaxis, :]
    input = torch.FloatTensor(input)
    input = input.to(torch.device("cuda"))
    output = model(input)

    props, preds = torch.max(output, 1)

    test_preds = preds.data.cpu().numpy()[0]

    return test_preds


if __name__ == '__main__':

    #########   set the GPU   ###########
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    #########  the input path   ##########
    import glob
    import tqdm
    import shutil

    bar = tqdm.tqdm()
    # resume = '/data1/sheng/temp/1211/2_output/weights-12-0-[0.9930313588850174]-1610505268.8762283.pth'
    resume = '/data1/sheng/1116/c/card_classify_output/weights-7-0-[0.9982909634693442]-1615519654.9759881.pth'
    model, test_transforms = class2_init(resume)

    right = 0
    error = 0
    img_list = glob.glob('/data1/sheng/1116/c/card_classify/*/*.jpg')
    # img_list = os.listdir(r'/data1/sheng/1116/04_geli_switch/train_dataset_output/test_error')
    for num, item in enumerate(img_list):
        # item = '/media/heils_lhl/data/youirobotData/light/2/1577264278.065616_ .jpg'

        img_ori = cv2.imread(item)
        img_ori = cv2.resize(img_ori, (48, 26))
        # img = cv2.cvtColor(img_ori,cv2.COLOR_BGR2RGB)
        img = img_ori
        testimage = [img]
        result = class2_forward(testimage,model,test_transforms)
        # print("result: ", result)
        # print(str(result[0][0]))

        gt = int(item.split('/')[-2])
        prd = int(result)
        # print(result,gt,item)
        if gt != prd:
            error += 1
            print('result: %d, gt: %d, item: %s' % (result, gt, item))
            org_name = item.split('/')[-1]
            # targretfolder = '/data1/sheng/1116/c/card_pred_error/'
            # if not os.path.exists(targretfolder):
            #     os.makedirs(targretfolder)
            # filename = item.split('/')[-1]
            filename = "gt_{}_pred_{}_{}.jpg".format(gt, prd, num)
            # cv2.imwrite(targretfolder+filename, img_ori)
            move_new_path = r'/data1/sheng/1116/c/second_dengdaifenlei/orggt_{}_'.format(gt) + org_name
            print("move ----  :", item, move_new_path)
            shutil.move(item, move_new_path)
            # exit(0)
            # print(targretfolder+filename)
        else:
            right += 1
        # cv2.imshow('wrong', img_ori)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        bar.update(1)
    print('right: %d, error: %d, acc: %s' % (right, error, right/(right+error)))